import base64
from os import path
import os
from flask import Flask
import io
from flask import Flask
from flask import send_file
import predict_DMQ_Axial
import matplotlib.pyplot as plt
import numpy as np
import gradcam
from PIL import Image


app = Flask(__name__)
app.config['DEBUG'] = True

# 返回一张原始图像和识别结果合并的图片
@app.route('/<id>')
def get_image_by_id(id):
    data = predict_DMQ_Axial.get_image(id)
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res
# 返回json格式的识别结果
@app.route('/result/<id>')
def get_result_by_id(id):
    data = predict_DMQ_Axial.main(id)
    return data

# 返回一张热力图和识别结果合并的图片
# @app.route('/heatmap/<id>')
def get_heatmap_by_id(id):
    data = gradcam.get_cam_img(id)
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res
@app.route('/heatmap/<filepath>')
def get_heatmap_by_filepath(filepath):
    b64 = base64.urlsafe_b64decode(filepath)
    path = bytes.decode(b64) 
    data = gradcam.get_cam_img(path)
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res
@app.route('/url/<filepath>')
def get_url_by_filepath(filepath):

    b64 = base64.urlsafe_b64decode(filepath)
    path = bytes.decode(b64) 
    return path
# 获取原始图像
@app.route('/rawimg/<id>')
def get_rawimg_by_id(id):
    data = gradcam.get_raw_img(id)
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res
# 存储并返回一张热力图和识别结果合并的图片
@app.route('/genheatmap/<id>')
def genarate_heatmap_and_get_from_folder(id):
    path='./save_img/'+ id + '/cam.jpg'
    if not os.path.exists(path):  
        gradcam.main(id)
    img = Image.open(path)
    fig = plt.figure()
    plt.axis('off')
    plt.imshow(img)
    canvas = fig.canvas
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()
    res = app.make_response(data)
    res.headers["Content-Type"] = "image/png"
    return res



if __name__ == '__main__':
    app.run()



